package main

import (
	"context"
	"encoding/json"
	"log"
	"net/http"
	"sync"
	"time"

	"github.com/go-redis/redis/v8"
	"github.com/gorilla/websocket"
)

var upgrader = websocket.Upgrader{
	ReadBufferSize:  1024,
	WriteBufferSize: 1024,
	CheckOrigin: func(r *http.Request) bool {
		return true
	},
}

const serverPassword = "myServerPassword"
const maxConnections = 4000 // 最大连接数
const redisQueueKey = "message_queue"

type Client struct {
	conn     *websocket.Conn
	userCode string
	key      string
	send     chan Message
}

type Message struct {
	UserCode       string      `json:"user_code"`
	Key            string      `json:"user_key"`
	ServerPassword string      `json:"server_password"`
	Message        interface{} `json:"message,omitempty"`
}

type Hub struct {
	clients     map[*Client]bool
	broadcast   chan Message
	register    chan *Client
	unregister  chan *Client
	users       map[string]*Client
	mu          sync.RWMutex
	count       int
	maxConn     int
	redisCtx    context.Context
	cancelFunc  context.CancelFunc
	redisClient *redis.Client
	isRunning   bool
	muIsRunning sync.Mutex
}

func newHub(maxConn int, redisClient *redis.Client) *Hub {
	redisCtx, cancelFunc := context.WithCancel(context.Background())
	return &Hub{
		broadcast:   make(chan Message),
		register:    make(chan *Client),
		unregister:  make(chan *Client),
		clients:     make(map[*Client]bool),
		users:       make(map[string]*Client),
		maxConn:     maxConn,
		redisCtx:    redisCtx,
		cancelFunc:  cancelFunc,
		redisClient: redisClient,
		isRunning:   false,
	}
}

func (h *Hub) run() {
	for {
		select {
		case client := <-h.register:
			h.mu.Lock()
			if h.count >= h.maxConn {
				log.Printf("已达到最大连接数，正在关闭用户的连接 %s", client.userCode)
				errMsg := Message{Message: "已达到最大连接数"}
				client.conn.WriteJSON(errMsg)
				client.conn.Close()
				h.mu.Unlock()
				continue
			}
			h.clients[client] = true
			h.users[client.userCode] = client
			h.count++
			h.mu.Unlock()
		case client := <-h.unregister:
			h.mu.Lock()
			if _, ok := h.clients[client]; ok {
				delete(h.clients, client)
				delete(h.users, client.userCode)
				close(client.send)
				h.count--
			}
			h.mu.Unlock()
		case message := <-h.broadcast:
			h.mu.RLock()
			messageToSend := message
			messageToSend.ServerPassword = "" // 清除 serverPassword 字段
			if message.UserCode == "" {
				for client := range h.clients {
					select {
					case client.send <- messageToSend:
					default:
						close(client.send)
						delete(h.clients, client)
						delete(h.users, client.userCode)
					}
				}
			} else if c, ok := h.users[message.UserCode]; ok {
				if c.key != message.Key || message.ServerPassword != serverPassword {
					log.Printf("用户身份验证失败 %s", message.UserCode)
					continue
				}
				select {
				case c.send <- messageToSend:
				default:
					close(c.send)
					delete(h.clients, c)
					delete(h.users, c.userCode)
				}
			}
			h.mu.RUnlock()
		}
	}
}

func serveWs(hub *Hub, w http.ResponseWriter, r *http.Request) {
	conn, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		log.Println(err)
		return
	}
	var msg Message
	err = conn.ReadJSON(&msg)
	if err != nil {
		log.Printf("错误: %v", err)
		conn.Close()
		return
	}
	if msg.ServerPassword != serverPassword {
		errMsg := Message{Message: "服务器密码无效"}
		conn.WriteJSON(errMsg)
		conn.Close()
		return
	}
	client := &Client{conn: conn, userCode: msg.UserCode, key: msg.Key, send: make(chan Message, 256)}
	hub.register <- client

	go client.writePump()
	go client.readPump()

	successMsg := Message{Message: "登录成功"}
	client.conn.WriteJSON(successMsg)
}

func (c *Client) readPump() {
	defer func() {
		hub.unregister <- c
		c.conn.Close()
	}()
	c.conn.SetReadLimit(512)
	c.conn.SetReadDeadline(time.Now().Add(60 * time.Second))
	c.conn.SetPongHandler(func(string) error { c.conn.SetReadDeadline(time.Now().Add(60 * time.Second)); return nil })
	for {
		var msg Message
		err := c.conn.ReadJSON(&msg)
		if err != nil {
			break
		}
		hub.broadcast <- msg
	}
}

func (c *Client) writePump() {
	ticker := time.NewTicker(10 * time.Second)
	defer func() {
		ticker.Stop()
		c.conn.Close()
	}()
	for {
		select {
		case message, ok := <-c.send:
			c.conn.SetWriteDeadline(time.Now().Add(10 * time.Second))
			if !ok {
				c.conn.WriteMessage(websocket.CloseMessage, []byte{})
				return
			}
			w, err := c.conn.NextWriter(websocket.TextMessage)
			if err != nil {
				return
			}
			err = json.NewEncoder(w).Encode(message)
			if err != nil {
				return
			}
			n := len(c.send)
			for i := 0; i < n; i++ {
				err = json.NewEncoder(w).Encode(<-c.send)
				if err != nil {
					return
				}
			}
			if err := w.Close(); err != nil {
				return
			}
		case <-ticker.C:
			c.conn.SetWriteDeadline(time.Now().Add(10 * time.Second))
			if err := c.conn.WriteMessage(websocket.PingMessage, nil); err != nil {
				return
			}
		}
	}
}

func sendMessageHandler(w http.ResponseWriter, r *http.Request) {
	var msg Message
	err := json.NewDecoder(r.Body).Decode(&msg)
	if err != nil {
		http.Error(w, err.Error(), http.StatusBadRequest)
		return
	}
	if msg.ServerPassword != serverPassword {
		http.Error(w, "服务器密码无效", http.StatusUnauthorized)
		return
	}
	if msg.UserCode != "" {
		hub.mu.RLock()
		c, ok := hub.users[msg.UserCode]
		hub.mu.RUnlock()
		if !ok {
			http.Error(w, "未找到用户", http.StatusNotFound)
			return
		}
		if c.key != msg.Key {
			http.Error(w, "用户密钥无效", http.StatusUnauthorized)
			return
		}
	}
	hub.broadcast <- msg
	w.WriteHeader(http.StatusOK)
	json.NewEncoder(w).Encode("发送成功")
}

func getUserCountHandler(w http.ResponseWriter, r *http.Request) {
	hub.mu.RLock()
	defer hub.mu.RUnlock()
	count := len(hub.clients)
	userCodes := make([]string, 0, count)
	for userCode := range hub.users {
		userCodes = append(userCodes, userCode)
	}
	response := struct {
		Count     int      `json:"count"`
		UserCodes []string `json:"user_codes"`
	}{
		Count:     count,
		UserCodes: userCodes,
	}
	json.NewEncoder(w).Encode(response)
}

func startRedisConsumer(hub *Hub) {
	hub.muIsRunning.Lock()
	hub.isRunning = true
	hub.muIsRunning.Unlock()
	for {
		select {
		case <-hub.redisCtx.Done():
			log.Println("停止Redis消费者")
			hub.muIsRunning.Lock()
			hub.isRunning = false
			hub.muIsRunning.Unlock()
			return
		default:
			result, err := hub.redisClient.BLPop(hub.redisCtx, 0, redisQueueKey).Result()
			if err != nil {
				if err == hub.redisCtx.Err() {
					log.Println("上下文已取消，Redis消费者停止")
					hub.muIsRunning.Lock()
					hub.isRunning = false
					hub.muIsRunning.Unlock()
					return
				}
				log.Printf("读取Redis队列时出错: %v", err)
				time.Sleep(1 * time.Second)
				continue
			}
			if len(result) > 1 {
				var msg Message
				err := json.Unmarshal([]byte(result[1]), &msg)
				if err != nil {
					log.Printf("解组Redis消息时出错: %v", err)
					continue
				}
				hub.broadcast <- msg
				log.Printf("已向用户发送消息 %s: %v", msg.UserCode, msg.Message)
			}
		}
	}
}

func controlRedisConsumerHandler(w http.ResponseWriter, r *http.Request) {
	action := r.URL.Query().Get("action")
	hub.muIsRunning.Lock()
	isRunning := hub.isRunning
	hub.muIsRunning.Unlock()

	switch action {
	case "start":
		if isRunning {
			http.Error(w, "Redis消费者已在运行", http.StatusBadRequest)
			return
		}
		if hub.cancelFunc == nil {
			hub.redisCtx, hub.cancelFunc = context.WithCancel(context.Background())
		}
		go startRedisConsumer(hub)
		w.WriteHeader(http.StatusOK)
		json.NewEncoder(w).Encode("Redis消费者已启动")
	case "stop":
		if !isRunning {
			http.Error(w, "Redis消费者已停止", http.StatusBadRequest)
			return
		}
		if hub.cancelFunc != nil {
			hub.cancelFunc()
			hub.cancelFunc = nil
		}
		w.WriteHeader(http.StatusOK)
		json.NewEncoder(w).Encode("Redis消费者已停止")
	default:
		http.Error(w, "无效的操作", http.StatusBadRequest)
	}
}

func getRedisConsumerStatusHandler(w http.ResponseWriter, r *http.Request) {
	hub.muIsRunning.Lock()
	status := hub.isRunning
	hub.muIsRunning.Unlock()
	response := struct {
		Status bool `json:"status"`
	}{
		Status: status,
	}
	json.NewEncoder(w).Encode(response)
}

var hub *Hub

func main() {
	redisClient := redis.NewClient(&redis.Options{
		Addr:     "localhost:6379",
		Password: "", // no password set
		DB:       0,  // use default DB
	})

	hub = newHub(maxConnections, redisClient)

	go hub.run()

	http.HandleFunc("/ws", func(w http.ResponseWriter, r *http.Request) {
		serveWs(hub, w, r)
	})
	http.HandleFunc("/send-message", sendMessageHandler)
	http.HandleFunc("/user-count", getUserCountHandler)
	http.HandleFunc("/control-consumer", controlRedisConsumerHandler)
	http.HandleFunc("/consumer-status", getRedisConsumerStatusHandler)

	log.Fatal(http.ListenAndServe(":8080", nil))
}
